"""
Copyright (C) University of Science and Technology of China.
Licensed under the MIT License.
"""

import sys
import argparse
import os
from util import util
import torch
import models
import data
import pickle


class BaseOptions():
    def __init__(self):
        self.initialized = False

    def initialize(self, parser):
        # experiment specifics
        parser.add_argument('--name', type=str, default='MichiGAN', help='name of the experiment. It decides where to store samples and models')

        parser.add_argument('--gpu_ids', type=str, default='0,1,2,3,4,5,6,7', help='gpu ids: e.g. 0  0,1,2, 0,2. use -1 for CPU')
        parser.add_argument('--checkpoints_dir', type=str, default='./checkpoints', help='models are saved here')
        parser.add_argument('--model', type=str, default='pix2pix', help='which model to use')
        parser.add_argument('--norm_G', type=str, default='spectralinstance', help='instance normalization or batch normalization')
        parser.add_argument('--norm_D', type=str, default='spectralinstance', help='instance normalization or batch normalization')
        parser.add_argument('--norm_E', type=str, default='spectralinstance', help='instance normalization or batch normalization')
        parser.add_argument('--phase', type=str, default='train', help='train, val, test, etc')
        parser.add_argument('--weight_norm_G', action='store_true', help='if specified, use weight normalization to replace feature norm in spade.')
        parser.add_argument('--weight_norm_g', type=int, default=0, help='0 means use the function by Pytorch, 1 means use ours.')
        # input/output sizes
        parser.add_argument('--batchSize', type=int, default=32, help='input batch size')
        parser.add_argument('--preprocess_mode', type=str, default='scale_width_and_crop', help='scaling and cropping of images at load time.', choices=("resize_and_crop", "crop", "scale_width", "scale_width_and_crop", "scale_shortside", "scale_shortside_and_crop", "fixed", "none"))
        parser.add_argument('--load_size', type=int, default=512, help='Scale images to this size. The final image will be cropped to --crop_size.')
        parser.add_argument('--crop_size', type=int, default=512, help='Crop to the width of crop_size (after initially scaling the images to load_size.)')
        parser.add_argument('--aspect_ratio', type=float, default=1.0, help='The ratio width/height. The final height of the load image will be crop_size/aspect_ratio')
        parser.add_argument('--label_nc', type=int, default=2, help='# of input label classes without unknown class. If you have unknown class as class label, specify --contain_dopntcare_label.')
        parser.add_argument('--contain_dontcare_label', action='store_true', help='if the label map contains dontcare label (dontcare=255)')
        parser.add_argument('--output_nc', type=int, default=3, help='# of output image channels')
        parser.add_argument('--orient_nc', type=int, default=2, help='# of orientation map channels')
        parser.add_argument('--add_noise_to_image', action='store_true', help='if specified, add noise to the image that remove hair')
        parser.add_argument('--use_original_image', action='store_true', help='if specified, use real image to generator')
        # parser.add_argument('--only_tag', action='store_true', help='if specified, not reference input')

        # for setting inputs
        parser.add_argument('--data_dir', type=str, default='/mnt/lvdisk1/tzt/HairSynthesis/SPADE-master/datasets/FFHQ',
                            help='path to the directory that contains training & val data')
        parser.add_argument('--dataroot', type=str, default='./datasets/cityscapes/')
        parser.add_argument('--dataset_mode', type=str, default='custom')
        parser.add_argument('--serial_batches', action='store_true', help='if true, takes images in order to make batches, otherwise takes them randomly')
        parser.add_argument('--no_flip', action='store_true', help='if specified, do not flip the images for data argumentation')
        parser.add_argument('--nThreads', default=8, type=int, help='# threads for loading data')
        parser.add_argument('--max_dataset_size', type=int, default=sys.maxsize, help='Maximum number of samples allowed per dataset. If the dataset directory contains more than max_dataset_size, only a subset is loaded.')
        parser.add_argument('--load_from_opt_file', action='store_true', help='load the options from checkpoints and use that as default')
        parser.add_argument('--cache_filelist_write', action='store_true', help='saves the current filelist into a text file, so that it loads faster')
        parser.add_argument('--cache_filelist_read', action='store_true', help='reads from the file list cache')
        parser.add_argument('--color_jitter', action='store_true', help='if specified, use color jitter to ref image.')
        parser.add_argument('--orient_random_disturb', action='store_true', help='if specified, random disturb the edges of orient for netG')
        parser.add_argument('--hair_random_disturb', action='store_true', help='if specified, random disturb the edges of hair for blend')

        # for displays
        parser.add_argument('--display_winsize', type=int, default=512, help='display window size')

        # for generator
        parser.add_argument('--netG', type=str, default='spadeb', help='selects model to use for netG (pix2pixhd | spade | spadeb)')
        parser.add_argument('--ngf', type=int, default=64, help='# of gen filters in first conv layer')
        parser.add_argument('--init_type', type=str, default='xavier', help='network initialization [normal|xavier|kaiming|orthogonal]')
        parser.add_argument('--init_variance', type=float, default=0.02, help='variance of the initialization distribution')
        parser.add_argument('--z_dim', type=int, default=256,
                            help="dimension of the latent z vector")
        parser.add_argument('--netIG', type=str, default='inpaint')
        parser.add_argument('--use_ig', action='store_true', help='which use inpainting generator to inpaint orientation')
        parser.add_argument('--ig_model_name', type=str, default='InpaintingModel_gen.pth', help='pretrained inpainting generator model')
        parser.add_argument('--norm_model', type=str, default='instance', help='normalization model [Batch | instance] in spaderesidualunet')
        parser.add_argument('--fix_netG', action='store_true',
                            help='if specified, do not update the weight of generator.')
        parser.add_argument('--num_upsampling_layers',
                            choices=('normal', 'more', 'most'), default='more',
                            help="If 'more', adds upsampling layer between the two middle resnet blocks. If 'most', also add one more upsampling + resnet layer at the end of the generator")
        parser.add_argument('--ms_step', type=int, default=0, help='The resolution of input for netG, [0,1,2,3,4].')
        parser.add_argument('--batch_sizes', type=str, default='32,32,32,16,8', help='The batch sizes for progressive training.')
        parser.add_argument('--alpha_value', type=float, default=-1, help='the alpha for progressive training.')
        parser.add_argument('--show_feat_maps', action='store_true', help='if specified, save the feature maps from generator in ./artifacts/.')

        # for image feature encoder like pix2pixHD
        parser.add_argument('--use_instance_feat', action='store_true', help='if specified, use feature encoder')
        parser.add_argument('--feat_num', type=int, default=3, help='the feature channels of feature encoder')
        parser.add_argument('--feat_input_nc', type=int, default=3, help='the input channels of feature encoder')

        # for reference image feature encoder
        parser.add_argument('--use_encoder', action='store_true', help='enable training with an image encoder.')
        parser.add_argument('--Image_encoder_mode', type=str, default='partialconv', help='encoder network [normal|instance|partialconv]')
        parser.add_argument('--norm_ref_encode', type=str, default='instance', help='[instance|none], none means no norm in ref encoder.')
        parser.add_argument('--ref_global_pool', action='store_true', help='if specified, use global pool in the ref encode.')

        # for blend network
        parser.add_argument('--use_blender', action='store_true', help='use blender to hold background.')
        parser.add_argument('--netB', type=str, default='blend2', help='select model to use for netB [blend | blend2]')
        parser.add_argument('--only_blend', action='store_true', help='only use blend for training or testing')

        # for instance-wise features
        parser.add_argument('--no_instance', default=True, help='if specified, do *not* add instance map as input')
        parser.add_argument('--nef', type=int, default=16, help='# of encoder filters in the first conv layer')
        parser.add_argument('--use_vae', action='store_true', help='enable training with an image encoder.')

        # for background image input for spadeb
        parser.add_argument('--noise_background', action='store_true', help='if specified, do noise to target image hair.')
        parser.add_argument('--random_expand_mask', action='store_true', help='if specified, expand the hair mask in the background encode for train.')
        parser.add_argument('--random_expand_th', type=float, default=0.05, help='the threshold of random expanding.')
        parser.add_argument('--bf_direct_add', action='store_true', help='if specified, direct add background feature without mask.')
        parser.add_argument('--random_noise_background', action='store_true', help='if specified, add noise to the background encoder.')

        parser.add_argument('--no_orientation', default=False, help='if specified, do *not* add orientation map as input')

        # for stroke orient inpainting
        parser.add_argument('--use_stroke', action='store_true', help='if specified, use stroke inpainting network.')
        parser.add_argument('--inpaint_mode', type=str, default='ref', choices=['ref', 'stroke'],
                            help='point out which inpaint network is used.')
        parser.add_argument('--netSIG', type=str, default='sinpaint')
        parser.add_argument('--sig_model_name', type=str, default='SInpaintingModel_gen.pth',
                            help='pretrained stroke inpainting generator model.')

        # for image padding zeros
        parser.add_argument('--add_zeros', action='store_true', help='if specified, add zeros to input data.')
        parser.add_argument('--add_feat_zeros', action='store_true', help='if specified, add zeros to the tensor before netG.')
        parser.add_argument('--add_th', type=int, default=64, help='total threshold of zeros padding.')
        # for clip the features
        parser.add_argument('--clip_th', type=float, default=300, help='the clip threshold for generator features.')
        parser.add_argument('--use_clip', action='store_true', help='if specified, clip the features in generator.')

        self.initialized = True
        return parser

    def gather_options(self):
        # initialize parser with basic options
        if not self.initialized:
            parser = argparse.ArgumentParser(
                formatter_class=argparse.ArgumentDefaultsHelpFormatter)
            parser = self.initialize(parser)

        # get the basic options
        opt, unknown = parser.parse_known_args()

        # modify model-related parser options
        model_name = opt.model
        model_option_setter = models.get_option_setter(model_name)
        parser = model_option_setter(parser, self.isTrain)

        # modify dataset-related parser options
        dataset_mode = opt.dataset_mode
        dataset_option_setter = data.get_option_setter(dataset_mode)
        parser = dataset_option_setter(parser, self.isTrain)

        opt, unknown = parser.parse_known_args()

        # if there is opt_file, load it.
        # The previous default options will be overwritten
        if opt.load_from_opt_file:
            parser = self.update_options_from_file(parser, opt)

        opt = parser.parse_args()
        self.parser = parser
        return opt

    def print_options(self, opt):
        message = ''
        message += '----------------- Options ---------------\n'
        for k, v in sorted(vars(opt).items()):
            comment = ''
            default = self.parser.get_default(k)
            if v != default:
                comment = '\t[default: %s]' % str(default)
            message += '{:>25}: {:<30}{}\n'.format(str(k), str(v), comment)
        message += '----------------- End -------------------'
        print(message)

    def option_file_path(self, opt, makedir=False):
        expr_dir = os.path.join(opt.checkpoints_dir, opt.name)
        if makedir:
            util.mkdirs(expr_dir)
        file_name = os.path.join(expr_dir, 'opt')
        return file_name

    def save_options(self, opt):
        file_name = self.option_file_path(opt, makedir=True)
        with open(file_name + '.txt', 'wt') as opt_file:
            for k, v in sorted(vars(opt).items()):
                comment = ''
                default = self.parser.get_default(k)
                if v != default:
                    comment = '\t[default: %s]' % str(default)
                opt_file.write('{:>25}: {:<30}{}\n'.format(str(k), str(v), comment))

        with open(file_name + '.pkl', 'wb') as opt_file:
            pickle.dump(opt, opt_file)

    def update_options_from_file(self, parser, opt):
        new_opt = self.load_options(opt)
        for k, v in sorted(vars(opt).items()):
            if hasattr(new_opt, k) and v != getattr(new_opt, k):
                new_val = getattr(new_opt, k)
                parser.set_defaults(**{k: new_val})
        return parser

    def load_options(self, opt):
        file_name = self.option_file_path(opt, makedir=False)
        new_opt = pickle.load(open(file_name + '.pkl', 'rb'))
        return new_opt

    def parse(self, save=False):

        opt = self.gather_options()
        opt.isTrain = self.isTrain   # train or test

        self.print_options(opt)
        if opt.isTrain:
            self.save_options(opt)

        # Set semantic_nc based on the option.
        # This will be convenient in many places
        opt.semantic_nc = opt.label_nc + \
            (1 if opt.contain_dontcare_label else 0) + \
            (0 if opt.no_instance else 1)

        # set gpu ids
        str_ids = opt.gpu_ids.split(',')
        opt.gpu_ids = []
        for str_id in str_ids:
            id = int(str_id)
            if id >= 0:
                opt.gpu_ids.append(id)
        if len(opt.gpu_ids) > 0:
            torch.cuda.set_device(opt.gpu_ids[0])

        assert len(opt.gpu_ids) == 0 or opt.batchSize % len(opt.gpu_ids) == 0, \
            "Batch size %d is wrong. It must be a multiple of # GPUs %d." \
            % (opt.batchSize, len(opt.gpu_ids))

        self.opt = opt
        return self.opt
